import argparse
import os
from pathlib import Path

import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import sys

sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor

sam_checkpoint = "C:\\Users\\crxc\\PycharmProjects\\segment-anything\\notebooks\\sam_vit_h_4b8939.pth"
model_type = "vit_h"

device = "cuda"


def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)


def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))


from PIL import Image
import numpy as np


def save_mask(mask, file_path):
    # 确保掩膜是2D的
    if mask.ndim == 3:
        mask = mask[:, :, 0]

    # 将布尔数组转换为0和255的uint8数组
    mask = (mask.astype(np.uint8) * 255)

    # 创建一个Pillow图像对象
    mask_image = Image.fromarray(mask)

    # 保存掩膜为PNG图像
    mask_image.save(file_path)


def generate_mask_singleimage(file_path, out_path1, out_path2, out_path3):
    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    sam.to(device=device)

    predictor = SamPredictor(sam)
    input_point = np.array([[2500, 2500], [2500, 600], [2500, 4800]])
    input_label = np.array([1, 0, 0])

    # input_point2 = np.array(
    #     [[1200, 600], [3000, 600], [2500, 2500], [1000, 5000], [3000, 5000], [400, 2500], [4200, 2500]])
    # input_label2 = np.array([0, 0, 1, 0, 0, 0, 0])

    image = cv2.imread(file_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    predictor.set_image(image)
    mask, score, logits = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        multimask_output=True,
    )

    # for i, (mask, score) in enumerate(zip(masks, scores)):
    #     plt.figure(figsize=(10, 10))
    #     plt.imshow(image)
    #     show_mask(mask, plt.gca())
    #     show_points(input_point, input_label, plt.gca())
    #     plt.title(f"Mask {i + 1}, Score: {score:.3f}", fontsize=18)
    #     plt.axis('off')
    #     plt.show()

    save_mask(mask[np.argmax(score)], out_path1)
    # save_mask(mask[1], getfilename(out_path2, score[1]))
    # save_mask(mask[2], getfilename(out_path3, score[2]))
    # mask2, score2, logits2 = predictor.predict(
    #     point_coords=input_point,
    #     point_labels=input_label2,
    #     multimask_output=False,
    # )
    # save_mask(mask2[0], getfilename(out_path2,score2[0]))


def getfilename(path, score):
    return path[:-4] + str(round(score, 2)) + ".png"


def generate_mask(in_dir, out_dir):
    # 使用os.listdir获取文件夹下所有文件和文件夹的名字
    for file_name in os.listdir(in_dir):
        if file_name.endswith('.png'):
            file_path = os.path.join(in_dir, file_name)
            out_dir1 = out_dir + "/" + "out1"
            checkPath(out_dir1)
            out_dir2 = out_dir + "/" + "out2"
            checkPath(out_dir2)
            out_dir3 = out_dir + "/" + "out3"
            checkPath(out_dir3)
            out_path1 = out_dir1 + "/" + file_name[:-4] + "mask.png"
            out_path2 = out_dir2 + "/" + file_name[:-4] + "mask.png"
            out_path3 = out_dir3 + "/" + file_name[:-4] + "mask.png"
            generate_mask_singleimage(file_path, out_path1, out_path2, out_path3)


def checkPath(out_dir1):
    folder_path = Path(out_dir1)
    if not folder_path.exists():
        folder_path.mkdir(parents=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='This is a large scale electron microscope data alignment scripy.')
    parser.add_argument("-i", "--in_dir", type=str, default="E:/wafer24/images", help="input dir")
    parser.add_argument("-o", "--out_dir", type=str, default="E:/wafer24/masks", help="output dir")
    args = parser.parse_args()
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    generate_mask(args.in_dir, args.out_dir)
